package com.pestdetection;

import org.datavec.api.io.labels.ParentPathLabelGenerator;
import org.datavec.api.split.FileSplit;
import org.datavec.image.loader.NativeImageLoader;
import org.datavec.image.recordreader.ImageRecordReader;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.dataset.api.preprocessor.DataNormalization;
import org.nd4j.linalg.dataset.api.preprocessor.ImagePreProcessingScaler;

import java.io.File;
import java.util.Random;
import java.awt.image.BufferedImage;
import javax.imageio.ImageIO;
import java.awt.Graphics2D;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.awt.RenderingHints;
import java.util.HashMap;
import java.util.Map;

/**
 * 数据加载器类
 * 负责加载和预处理病虫害图片数据
 * 主要功能：
 * 1. 调整图片大小
 * 2. 预处理图片
 * 3. 验证数据集
 * 4. 加载训练数据
 */
public class DataLoader {
    // 图片处理参数
    private static final int 图片通道数 = 3;    // RGB彩色图片
    private static final int 批处理大小 = 1;    // 每次处理1张图片
    private static final int 目标高度 = 224;    // 目标图片高度
    private static final int 目标宽度 = 224;    // 目标图片宽度
    private static final Random 随机数生成器 = new Random(123); // 用于打乱数据顺序

    /**
     * 调整图片大小，保持宽高比
     * @param 原始图片 需要调整大小的图片
     * @return 调整后的图片
     */
    private static BufferedImage 调整图片大小(BufferedImage 原始图片) {
        int 原始宽度 = 原始图片.getWidth();
        int 原始高度 = 原始图片.getHeight();
        
        // 计算缩放比例
        double 宽度比例 = (double) 目标宽度 / 原始宽度;
        double 高度比例 = (double) 目标高度 / 原始高度;
        double 缩放比例 = Math.min(宽度比例, 高度比例);
        
        int 新宽度 = (int) (原始宽度 * 缩放比例);
        int 新高度 = (int) (原始高度 * 缩放比例);
        
        // 创建新图片
        BufferedImage 调整后的图片 = new BufferedImage(目标宽度, 目标高度, BufferedImage.TYPE_INT_RGB);
        Graphics2D 画笔 = 调整后的图片.createGraphics();
        
        // 设置图片质量
        画笔.setRenderingHint(RenderingHints.KEY_INTERPOLATION, RenderingHints.VALUE_INTERPOLATION_BILINEAR);
        画笔.setRenderingHint(RenderingHints.KEY_RENDERING, RenderingHints.VALUE_RENDER_QUALITY);
        画笔.setRenderingHint(RenderingHints.KEY_ANTIALIASING, RenderingHints.VALUE_ANTIALIAS_ON);
        
        // 计算居中位置
        int 水平位置 = (目标宽度 - 新宽度) / 2;
        int 垂直位置 = (目标高度 - 新高度) / 2;
        
        // 绘制图片
        画笔.drawImage(原始图片, 水平位置, 垂直位置, 新宽度, 新高度, null);
        画笔.dispose();
        
        return 调整后的图片;
    }

    /**
     * 预处理图片
     * @param 图片文件 需要处理的图片文件
     * @return 处理后的图片
     * @throws IOException 如果处理失败
     */
    private static BufferedImage 预处理图片(File 图片文件) throws IOException {
        BufferedImage 原始图片 = ImageIO.read(图片文件);
        if (原始图片 == null) {
            throw new IOException("无法读取图片: " + 图片文件.getAbsolutePath());
        }
        
        // 检查图片尺寸
        if (原始图片.getWidth() != 目标宽度 || 原始图片.getHeight() != 目标高度) {
            System.out.println("调整图片尺寸: " + 图片文件.getName() + 
                             " (" + 原始图片.getWidth() + "x" + 原始图片.getHeight() + 
                             " -> " + 目标宽度 + "x" + 目标高度 + ")");
            return 调整图片大小(原始图片);
        }
        
        return 原始图片;
    }

    /**
     * 验证数据集
     * @param 数据集路径 数据集所在路径
     * @return 病虫害种类列表
     * @throws IOException 如果数据集无效
     */
    private static List<String> 验证数据集(String 数据集路径) throws IOException {
        File 数据集目录 = new File(数据集路径);
        if (!数据集目录.exists()) {
            throw new IOException("数据集目录不存在: " + 数据集路径);
        }

        File[] 种类目录 = 数据集目录.listFiles(File::isDirectory);
        if (种类目录 == null || 种类目录.length == 0) {
            throw new IOException("数据集目录为空或没有子目录: " + 数据集路径);
        }

        Map<String, Integer> 种类图片数量 = new HashMap<>();
        List<String> 有效图片 = new ArrayList<>();
        List<String> 种类名称 = new ArrayList<>();
        
        for (File 种类文件夹 : 种类目录) {
            String 种类名 = 种类文件夹.getName();
            种类名称.add(种类名);
            int 图片数量 = 0;
            
            File[] 图片文件 = 种类文件夹.listFiles((dir, name) -> {
                String 小写文件名 = name.toLowerCase();
                return 小写文件名.endsWith(".jpg") || 小写文件名.endsWith(".png");
            });
            
            if (图片文件 != null) {
                for (File 图片 : 图片文件) {
                    try {
                        预处理图片(图片);
                        有效图片.add(图片.getAbsolutePath());
                        图片数量++;
                    } catch (IOException e) {
                        System.err.println("警告: 跳过无效图片 " + 图片.getAbsolutePath() + 
                                         ": " + e.getMessage());
                    }
                }
            }
            
            种类图片数量.put(种类名, 图片数量);
        }

        if (有效图片.isEmpty()) {
            throw new IOException("数据集中没有有效的图片文件");
        }

        // 打印数据集统计信息
        System.out.println("数据集验证通过，共发现 " + 有效图片.size() + " 张有效图片，分布在 " + 
                          种类图片数量.size() + " 个种类中");
        
        // 打印每个种类的图片数量和标签映射
        System.out.println("\n种类分布情况：");
        int 标签索引 = 0;
        for (Map.Entry<String, Integer> 种类信息 : 种类图片数量.entrySet()) {
            System.out.println(String.format("- %s: %d 张图片 (标签索引: %d)", 
                种类信息.getKey(), 种类信息.getValue(), 标签索引++));
        }

        return 种类名称;
    }

    /**
     * 加载训练数据
     * @param 数据集路径 数据集所在路径
     * @return 处理后的数据集
     * @throws Exception 处理过程中可能出现的异常
     */
    public static DataSetIterator loadData(String 数据集路径) throws Exception {
        try {
            // 验证数据集并获取种类列表
            List<String> 种类列表 = 验证数据集(数据集路径);
            int 种类数量 = 种类列表.size();

            File 数据集目录 = new File(数据集路径);
            FileSplit 文件分割 = new FileSplit(数据集目录, NativeImageLoader.ALLOWED_FORMATS, 随机数生成器);
            
            // 使用文件夹名称作为标签
            ParentPathLabelGenerator 标签生成器 = new ParentPathLabelGenerator();
            
            // 创建图片读取器
            ImageRecordReader 图片读取器 = new ImageRecordReader(目标高度, 目标宽度, 图片通道数, 标签生成器);
            图片读取器.initialize(文件分割);
            
            System.out.println("检测到 " + 种类数量 + " 个种类");
            
            // 创建数据集迭代器
            RecordReaderDataSetIterator 数据集 = new RecordReaderDataSetIterator.Builder(图片读取器, 批处理大小)
                .classification(-1, 种类数量)  // 使用-1表示从metadata获取标签
                .maxNumBatches(10)  // 限制批次数量以便调试
                .preProcessor(null)  // 暂时禁用预处理器
                .build();
            
            // 禁用异步预取和元数据收集
            数据集.setCollectMetaData(false);
            
            // 数据标准化
            DataNormalization 标准化器 = new ImagePreProcessingScaler(0, 1);
            标准化器.fit(数据集);
            数据集.setPreProcessor(标准化器);
            
            // 验证标签维度
            if (数据集.getLabels().size() != 种类数量) {
                throw new IllegalStateException(String.format(
                    "标签维度不匹配：期望 %d 个种类，实际检测到 %d 个种类",
                    种类数量, 数据集.getLabels().size()));
            }
            
            return 数据集;
            
        } catch (Exception e) {
            System.err.println("数据加载错误：" + e.getMessage());
            e.printStackTrace();
            throw e;
        }
    }
} 